# Obtained from: https://github.com/vikolss/DACS
# Copyright (c) 2020 vikolss. Licensed under the MIT License
# A copy of the license is available at resources/license_dacs

import kornia
import numpy as np
import torch
import torch.nn as nn

# 1.数据增强：
# 混合操作（one_mix）：根据掩码对输入的图像和标签进行线性混合。
# 颜色抖动（color_jitter）：随机调整图像的颜色属性（亮度、对比度、饱和度、色调），以增加数据多样性。
# 高斯模糊（gaussian_blur）：对图像应用高斯模糊，模拟不同条件下的图像质量变化。
# 2.归一化与去归一化：
# 去归一化（denorm, denorm_）：将归一化后的图像数据恢复到原始范围，通常用于在应用增强操作前。
# 重新归一化（renorm_）：将图像数据重新归一化到指定的均值和标准差范围内，通常用于在应用增强操作后。
# 3.参数获取与处理：
# 获取均值和标准差（get_mean_std）：从图像元数据中提取均值和标准差，用于归一化和去归一化操作。
# 4.类别掩码生成：
# 生成类别掩码（get_class_masks, generate_class_mask）：根据标签生成类别掩码，用于后续的混合操作或其他处理。
# 5.综合增强函数：
# 综合增强（strong_transform）：结合上述多种增强方法，根据传入的参数对图像和标签进行一系列增强处理，确保数据的多样性和鲁棒性。

def strong_transform(param, data=None, target=None):
    """这里data是两个域图像的堆叠，target是两个域标签的堆叠"""
    assert ((data is not None) or (target is not None))

    data, target = one_mix(mask=param['mix'], data=data, target=target)

    data, target = color_jitter(
        color_jitter=param['color_jitter'],
        s=param['color_jitter_s'],
        p=param['color_jitter_p'],
        mean=param['mean'],
        std=param['std'],
        data=data,
        target=target)
    data, target = gaussian_blur(blur=param['blur'], data=data, target=target)
    return data, target

# 获取均值和标准差
def get_mean_std(img_metas, dev):
    mean = [
        torch.as_tensor(img_metas[i]['img_norm_cfg']['mean'], device=dev)
        for i in range(len(img_metas))
    ]
    mean = torch.stack(mean).view(-1, 3, 1, 1)
    std = [
        torch.as_tensor(img_metas[i]['img_norm_cfg']['std'], device=dev)
        for i in range(len(img_metas))
    ]
    std = torch.stack(std).view(-1, 3, 1, 1)
    return mean, std


def denorm(img, mean, std):
    return img.mul(std).add(mean) / 255.0


def denorm_(img, mean, std):
    img.mul_(std).add_(mean).div_(255.0)


def renorm_(img, mean, std):
    img.mul_(255.0).sub_(mean).div_(std)


def color_jitter(color_jitter, mean, std, data=None, target=None, s=.25, p=.2):
    # s is the strength of colorjitter
    if not (data is None):
        if data.shape[1] == 3:
            if color_jitter > p:
                if isinstance(s, dict):
                    seq = nn.Sequential(kornia.augmentation.ColorJitter(**s))
                else:
                    seq = nn.Sequential(
                        kornia.augmentation.ColorJitter(
                            brightness=s, contrast=s, saturation=s, hue=s))
                denorm_(data, mean, std)
                data = seq(data)
                renorm_(data, mean, std)
    return data, target


def gaussian_blur(blur, data=None, target=None):
    if not (data is None):
        if data.shape[1] == 3:
            if blur > 0.5:
                sigma = np.random.uniform(0.15, 1.15)
                kernel_size_y = int(
                    np.floor(
                        np.ceil(0.1 * data.shape[2]) - 0.5 +
                        np.ceil(0.1 * data.shape[2]) % 2))
                kernel_size_x = int(
                    np.floor(
                        np.ceil(0.1 * data.shape[3]) - 0.5 +
                        np.ceil(0.1 * data.shape[3]) % 2))
                kernel_size = (kernel_size_y, kernel_size_x)
                seq = nn.Sequential(
                    kornia.filters.GaussianBlur2d(
                        kernel_size=kernel_size, sigma=(sigma, sigma)))
                data = seq(data)
    return data, target


def get_class_masks(labels):
    class_masks = []
    for label in labels:
        classes = torch.unique(labels)
        nclasses = classes.shape[0]
        class_choice = np.random.choice(
            nclasses, int((nclasses + nclasses % 2) / 2), replace=False)
        classes = classes[torch.Tensor(class_choice).long()]
        class_masks.append(generate_class_mask(label, classes).unsqueeze(0))
    return class_masks


def generate_class_mask(label, classes):
    label, classes = torch.broadcast_tensors(label,
                                             classes.unsqueeze(1).unsqueeze(2))
    class_mask = label.eq(classes).sum(0, keepdims=True)
    return class_mask


def one_mix(mask, data=None, target=None):
    """
    根据提供的mask，data和target，执行一种混合操作。

    参数:
    - mask: 一个包含掩码的元组，用于决定data和target的混合方式。
    - data: 一个包含要混合的数据的元组。
    - target: 一个包含要混合的目标的元组。

    返回:
    - data: 混合后的数据。
    - target: 混合后的目标。
    """
    # 如果mask为None，则直接返回原始的data和target，不做任何混合操作。
    if mask is None:
        return data, target

    # 如果data不为None，则执行数据的混合操作。
    if not (data is None):
        # 使用torch.broadcast_tensors对mask和data的第一个元素进行广播，以匹配它们的形状。
        stackedMask0, _ = torch.broadcast_tensors(mask[0], data[0])
        # 根据广播后的mask，按比例混合data的第一个和第二个元素。
        data = (stackedMask0 * data[0] +
                (1 - stackedMask0) * data[1]).unsqueeze(0)

    # 如果target不为None，则执行目标的混合操作。
    if not (target is None):
        # 使用torch.broadcast_tensors对mask和target的第一个元素进行广播，以匹配它们的形状。
        stackedMask0, _ = torch.broadcast_tensors(mask[0], target[0])
        # 根据广播后的mask，按比例混合target的第一个和第二个元素。
        target = (stackedMask0 * target[0] +
                  (1 - stackedMask0) * target[1]).unsqueeze(0)

    # 返回混合后的data和target。
    return data, target
